/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef GE_INFER_SHAPE_CONTEXT_H
#define GE_INFER_SHAPE_CONTEXT_H
#include "graph/shape.h"
#include "graph/tensor.h"
#include "register/runtime_context.h"
#include "register/runtime_attrs.h"

namespace ge {
class Node;

class GRAPH_API_EXPORT InferShapeContext : public RuntimeContext {
public:
    InferShapeContext(ge::Node& node);
    /**
     * 根据输入index，获取输入shape指针
     * @param index 输入index
     * @return 输入shape指针，index非法时，返回空指针
     */
    const Shape* GetInputShape(const size_t index) const;

    /**
     * 基于算子IR原型定义，获取`OPTIONAL_INPUT`类型的输入shape指针
     * @param ir_index IR原型定义中的index
     * @return shape指针，index非法，或该INPUT没有实例化时，返回空指针
     */
    const Shape* GetOptionalInputShape(const size_t ir_index) const;
    /**
     * 基于算子IR原型定义，获取`DYNAMIC_INPUT`类型的输入shape指针
     * @param ir_index IR原型定义中的index
     * @param relative_index
     * 该输入实例化后的相对index，例如某个DYNAMIC_INPUT实例化了3个输入，那么relative_index的有效范围是[0,2]
     * @return shape指针，index或relative_index非法时，返回空指针
     */
    const Shape* GetDynamicInputShape(const size_t ir_index, const size_t relative_index) const;

    /**
     * 根据输出index，获取输出shape指针
     * @param index 输出index
     * @return 输出shape指针，index非法时，返回空指针
     */
    Shape* GetOutputShape(const size_t index);
};

} // namespace ge
#endif // GE_INFER_SHAPE_CONTEXT_H
